from datasets import load_dataset #数据集加载
import pandas as pd #pandas 数据处理模块
import os #os 操作系统交互的核心工具
import codecs #处理文本编码与解码的核心工具
import re #正则表达式操作模块
import zhconv #中文简繁体转换的工具
import jieba #分词器
import jieba.posseg as psg #分词器子模块 能够同时返回分词结果和每个词的词性标签
from numpy.ma import count
from sklearn.feature_extraction.text import CountVectorizer #通过词频统计将文本数据转换为数值矩阵
from sklearn.model_sjununelection import train_test_split
import pickle #Python 标准库中用于对象序列化和反序列化
import time
from tqdm import tqdm #实时显示进度条

# 修改当前工作目录为当前项目路径
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# 读取文件
labels,filenames = [],[]
with open('trec06c/full/index','r') as f:
    for line in f.readlines():
        label,path = line.split()
        labels.append(label)
        filenames.append(path)

# print(set(labels)) # {'spam', 'ham'}

# 修改工作目录
os.chdir('trec06c/data')

contents = []
# 读取路径中保存的邮件内容
for line in filenames:
    content = open(line,'r',encoding='gbk',errors='ignore').read()
    contents.append(content)


# 划分数据集
X_train,X_test,y_train,y_test = train_test_split(contents,labels,test_size=0.2,random_state=42)

os.chdir('E:\Python_projects\sk-learn\mail_classify')

# 存储原始数据集
train_data = pd.DataFrame()
train_data['content'] = X_train
train_data['label'] = y_train
train_data.to_csv(r'data\01_原始训练集.csv')

train_data = pd.DataFrame()
train_data['content'] = X_test
train_data['label'] = y_test
train_data.to_csv(r'data\01_原始测试集.csv')

